In the notebook, the background data should be broadcasted. When the explain_instances dataframe (observations to be explained) is in a mid range (50 to 100-ish), Spark will use a unexpected type of join plan, and messes up with the parallelization of the Kernel SHAP sampler, thus creating a performance bottleneck. Broadcasting the background dataset makes Spark respect the partitioning of the explain_instances dataframe.

These two notebooks both explain only 5 data points, so the performance bottleneck is not obvious. However, if we change 5 to 50, it becomes obvious. But if we further change it 500, Spark uses the intended join plan, and the bottleneck is not triggered.

I thought about forcing the broadcast inside the explainer, but this may create unexpected effect for other scenarios, so I'm hesitant to do so.